# 生成一个持久化的、包含所有代码知识的向量数据库。
# 这个数据库就是AI的“长期记忆”，随时准备着被检索和调用。
import os
import sys
from dotenv import load_dotenv
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
# from langchain_openai import OpenAIEmbeddings
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma



# --- 1. 加载环境变量 ---
load_dotenv()

# 命令行信息
arguments = sys.argv
dir_name = arguments[1]


# --- 2. 定义常量 ---
REPO_PATH = f"../test_repo/{dir_name}"
PERSIST_DIRECTORY = "../chroma_db"  # 向量数据库的存储路径


def ingest_codebase():
    """
    加载、切分代码库，并将其嵌入到向量数据库中。
    """
    print("--- 开始加载代码库文件 ---")

    # 使用DirectoryLoader加载整个项目，并指定加载.py和.md文件
    # glob参数用于匹配文件模式
    loader = DirectoryLoader(
        REPO_PATH,
        glob="**/*.py",  # 只加载Python文件
        loader_cls=TextLoader,
        use_multithreading=True,
        show_progress=True,
        loader_kwargs={"autodetect_encoding": True}  # 自动检测编码
    )
    documents = loader.load()
    print(f"--- 成功加载 {len(documents)} 个文件 ---")

    # --- 3. 切分文档 ---
    # 对于代码，使用专门的切分器效果更好，这里为了简化使用通用的
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,  # 每个块的最大字符数
        chunk_overlap=200  # 块之间的重叠字符数
    )
    docs = text_splitter.split_documents(documents)
    print(f"--- 成功将文档切分为 {len(docs)} 个片段 ---")

    # --- 4. 创建嵌入并存入 ChromaDB ---
    print("--- 开始创建文本嵌入并存入向量数据库 ---")
    # 这个embeddings对象就是一个转换工具，你给它一段文本，它就给你返回一个代表其语义的向量。
    embeddings = OllamaEmbeddings(model="nomic-embed-text")

    # 从文档创建向量数据库，并持久化到磁盘
    vectordb = Chroma.from_documents(
        documents=docs,
        embedding=embeddings,
        persist_directory=PERSIST_DIRECTORY
    )
    vectordb.persist()
    print(f"--- 成功将数据存入 {PERSIST_DIRECTORY} ---")


if __name__ == "__main__":
    ingest_codebase()